#include "CodeGenerator.h"
#include <iostream>
#include <fstream>
#include <sstream>
#include <algorithm>
#include <unordered_map>
#include <set>

// ==================== RegisterAllocator 实现 ====================

RegisterAllocator::RegisterAllocator() {
    // 初始化所有寄存器为可用状态
    // 优先分配t寄存器（调用者保存）
    for (int i = static_cast<int>(Register::T0); i <= static_cast<int>(Register::T6); ++i) {
        freeRegisters.push(static_cast<Register>(i));
        registerUsage[static_cast<Register>(i)] = false;
    }
    // 然后分配a寄存器（参数寄存器）
    for (int i = static_cast<int>(Register::A0); i <= static_cast<int>(Register::A7); ++i) {
        freeRegisters.push(static_cast<Register>(i));
        registerUsage[static_cast<Register>(i)] = false;
    }
    // 最后分配s寄存器（被调用者保存，尽量避免使用）
    for (int i = static_cast<int>(Register::S0); i <= static_cast<int>(Register::S11); ++i) {
        freeRegisters.push(static_cast<Register>(i));
        registerUsage[static_cast<Register>(i)] = false;
    }
    
    // 初始化特殊寄存器
    registerUsage[Register::ZERO] = false;
    registerUsage[Register::RA] = false;
    registerUsage[Register::SP] = false;
    registerUsage[Register::GP] = false;
    registerUsage[Register::TP] = false;
}

Register RegisterAllocator::allocateRegister() {
    if (!freeRegisters.empty()) {
        Register reg = freeRegisters.top();
        freeRegisters.pop();
        registerUsage[reg] = true;
        return reg;
    }
    
    // 没有可用寄存器，需要溢出处理
    // 优先溢出临时寄存器
    for (int i = static_cast<int>(Register::T0); i <= static_cast<int>(Register::T6); ++i) {
        Register reg = static_cast<Register>(i);
        if (registerUsage[reg]) {
            return spillRegister(reg);
        }
    }
    
    // 如果临时寄存器都用完了，溢出参数寄存器
    for (int i = static_cast<int>(Register::A0); i <= static_cast<int>(Register::A7); ++i) {
        Register reg = static_cast<Register>(i);
        if (registerUsage[reg]) {
            return spillRegister(reg);
        }
    }
    
    // 最后才考虑溢出保存寄存器（尽量避免）
    for (int i = static_cast<int>(Register::S0); i <= static_cast<int>(Register::S11); ++i) {
        Register reg = static_cast<Register>(i);
        if (registerUsage[reg]) {
            return spillRegister(reg);
        }
    }
    
    // 如果所有寄存器都被使用，返回T0（这种情况不应该发生）
    return Register::T0;
}

void RegisterAllocator::freeRegister(Register reg) {
    if (reg == Register::ZERO) return;
    
    registerUsage[reg] = false;
    freeRegisters.push(reg);
}

Register RegisterAllocator::getVariableRegister(const std::string& varName) {
    auto it = variableRegisters.find(varName);
    if (it != variableRegisters.end()) {
        return it->second;
    }
    return Register::ZERO;
}

void RegisterAllocator::setVariableRegister(const std::string& varName, Register reg) {
    variableRegisters[varName] = reg;
}

void RegisterAllocator::clearVariableRegisters() {
    variableRegisters.clear();
}

std::string RegisterAllocator::registerToString(Register reg) {
    switch (reg) {
        case Register::ZERO: return "x0";
        case Register::RA: return "ra";
        case Register::SP: return "sp";
        case Register::GP: return "gp";
        case Register::TP: return "tp";
        case Register::T0: return "t0";
        case Register::T1: return "t1";
        case Register::T2: return "t2";
        case Register::S0: return "s0";
        case Register::S1: return "s1";
        case Register::A0: return "a0";
        case Register::A1: return "a1";
        case Register::A2: return "a2";
        case Register::A3: return "a3";
        case Register::A4: return "a4";
        case Register::A5: return "a5";
        case Register::A6: return "a6";
        case Register::A7: return "a7";
        case Register::S2: return "s2";
        case Register::S3: return "s3";
        case Register::S4: return "s4";
        case Register::S5: return "s5";
        case Register::S6: return "s6";
        case Register::S7: return "s7";
        case Register::S8: return "s8";
        case Register::S9: return "s9";
        case Register::S10: return "s10";
        case Register::S11: return "s11";
        case Register::T3: return "t3";
        case Register::T4: return "t4";
        case Register::T5: return "t5";
        case Register::T6: return "t6";
        default: return "unknown";
    }
}

// 新增：寄存器溢出管理
// RegisterAllocator::spillRegister 禁止对s寄存器做sp负偏移保存
Register RegisterAllocator::spillRegister(Register reg) {
    // 只允许溢出t寄存器和a寄存器，禁止s寄存器
    if (reg >= Register::S0 && reg <= Register::S11) {
        std::cerr << "[ERROR] 禁止对s寄存器溢出保存: " << registerToString(reg) << std::endl;
        return reg; // 直接返回，不做溢出
    }
    spilledRegisters[reg] = nextSpillOffset;
    nextSpillOffset -= 4;
    return reg;
}

// 新增：生成寄存器溢出指令
// generateSpillInstruction 禁止对s寄存器生成保存指令
void RegisterAllocator::generateSpillInstruction(Register reg, std::vector<std::string>& code) {
    if (reg >= Register::S0 && reg <= Register::S11) return;
    auto it = spilledRegisters.find(reg);
    if (it != spilledRegisters.end()) {
        code.push_back("    sw " + registerToString(reg) + ", " + std::to_string(it->second) + "(sp)");
    }
}

// 新增：生成寄存器恢复指令
// generateUnspillInstruction 禁止对s寄存器生成恢复指令
void RegisterAllocator::generateUnspillInstruction(Register reg, std::vector<std::string>& code) {
    if (reg >= Register::S0 && reg <= Register::S11) return;
    auto it = spilledRegisters.find(reg);
    if (it != spilledRegisters.end()) {
        code.push_back("    lw " + registerToString(reg) + ", " + std::to_string(it->second) + "(sp)");
    }
}

void RegisterAllocator::unspillRegister(Register reg) {
    auto it = spilledRegisters.find(reg);
    if (it != spilledRegisters.end()) {
        spilledRegisters.erase(it);
    }
}

bool RegisterAllocator::isCallerSaved(Register reg) {
    // t0-t6, a0-a7 是调用者保存寄存器
    return (reg >= Register::T0 && reg <= Register::T6) ||
           (reg >= Register::A0 && reg <= Register::A7);
}

bool RegisterAllocator::isCalleeSaved(Register reg) {
    // s0-s11 是被调用者保存寄存器
    return (reg >= Register::S0 && reg <= Register::S11);
}

std::vector<Register> RegisterAllocator::getActiveCallerSavedRegisters() {
    std::vector<Register> activeRegs;
    for (const auto& pair : registerUsage) {
        if (pair.second && isCallerSaved(pair.first)) {
            activeRegs.push_back(pair.first);
        }
    }
    return activeRegs;
}

void RegisterAllocator::allocateSpecificRegister(Register reg) {
    if (reg == Register::ZERO) return;
    registerUsage[reg] = true;
    // 从freeRegisters移除reg（如果存在）
    std::stack<Register> tmp;
    while (!freeRegisters.empty()) {
        Register top = freeRegisters.top();
        freeRegisters.pop();
        if (top != reg) tmp.push(top);
    }
    while (!tmp.empty()) {
        freeRegisters.push(tmp.top());
        tmp.pop();
    }
}

bool RegisterAllocator::hasFreeRegisters() const {
    return !freeRegisters.empty();
}

// ==================== CodeGenerator 实现 ====================

CodeGenerator::CodeGenerator() 
    : labelCounter(0), stackOffset(0) {
    emitComment("RISC-V32 Assembly Code Generated by Compiler");
    emit("");
}

std::string CodeGenerator::generateLabel(const std::string& prefix) {
    return prefix + "_" + std::to_string(labelCounter++);
}

// emit函数增加日志输出，辅助定位所有sw ... (sp)的来源
void CodeGenerator::emit(const std::string& instruction) {
    // 检查是否sw ... (sp)
    if (instruction.find("sw ") != std::string::npos && instruction.find("(sp)") != std::string::npos) {
        size_t regStart = instruction.find("sw ") + 3;
        size_t regEnd = instruction.find(", ");
        std::string reg = instruction.substr(regStart, regEnd - regStart);
        size_t offsetStart = regEnd + 2;
        size_t offsetEnd = instruction.find("(", offsetStart);
        int offset = std::stoi(instruction.substr(offsetStart, offsetEnd - offsetStart));
        std::cerr << "[emit-sw-sp] func: " << currentFunction << ", reg: " << reg << ", offset: " << offset << ", instr: " << instruction << std::endl;
    }
    code.push_back(instruction);
}

void CodeGenerator::emitComment(const std::string& comment) {
    code.push_back("# " + comment);
}

void CodeGenerator::emitFunctionPrologue(const std::string& funcName, int frameSize) {
    emit("");
    emit(funcName + ":");
    emitComment("Function prologue");
    emit("    addi sp, sp, -" + std::to_string(frameSize));
    emit("    sw ra, 0(sp)");
    emit("    sw s0, 4(sp)");
    emit("    sw s1, 8(sp)");
    emit("    sw s2, 12(sp)");
    emit("    sw s3, 16(sp)");
    emit("    sw s4, 20(sp)");
    emit("    sw s5, 24(sp)");
    emit("    sw s6, 28(sp)");
    emit("    sw s7, 32(sp)");
    emit("    sw s8, 36(sp)");
    emit("    sw s9, 40(sp)");
    emit("    sw s10, 44(sp)");
    emit("    sw s11, 48(sp)");
    emit("    addi s0, sp, " + std::to_string(frameSize)); // s0指向原来的sp位置

    // 参数仍然留在 a0-a7 中，不做额外转移；调用其它函数前由 saveRegisters 自动保存活动的 a 寄存器
    stackOffset = 0;
    nextOffset = 0;
    currentFrameSize = frameSize;
}

void CodeGenerator::emitFunctionEpilogue(int frameSize) {
    emitComment("Function epilogue");
    emit("    lw ra, 0(sp)");
    emit("    lw s0, 4(sp)");
    emit("    lw s1, 8(sp)");
    emit("    lw s2, 12(sp)");
    emit("    lw s3, 16(sp)");
    emit("    lw s4, 20(sp)");
    emit("    lw s5, 24(sp)");
    emit("    lw s6, 28(sp)");
    emit("    lw s7, 32(sp)");
    emit("    lw s8, 36(sp)");
    emit("    lw s9, 40(sp)");
    emit("    lw s10, 44(sp)");
    emit("    lw s11, 48(sp)");
    emit("    addi sp, sp, " + std::to_string(frameSize));
    emit("    jr ra");
}

void CodeGenerator::saveRegisters() {
    savedCallerRegisters.clear();
    currentCallerSaveStackOffset = 0;
    std::vector<Register> activeCallerSaved = regAlloc.getActiveCallerSavedRegisters();
    
    // 保存所有调用者保存寄存器：t0-t6 以及 a0-a7
    std::vector<Register> registersToSave;
    for (Register reg : activeCallerSaved) {
        if (regAlloc.isCallerSaved(reg)) {
            registersToSave.push_back(reg);
        }
    }
    
    if (registersToSave.empty()) return;
    
    int bytesToSave = registersToSave.size() * 4;
    emit("    addi sp, sp, -" + std::to_string(bytesToSave));
    int offset = 0;
    for (Register reg : registersToSave) {
        emit("    sw " + regAlloc.registerToString(reg) + ", " + std::to_string(offset) + "(sp)");
        savedCallerRegisters.push_back({reg, offset});
        offset += 4;
    }
}

void CodeGenerator::restoreRegisters() {
    if (savedCallerRegisters.empty()) return;
    for (auto it = savedCallerRegisters.rbegin(); it != savedCallerRegisters.rend(); ++it) {
        Register reg = it->first;
        int offset = it->second;
        emit("    lw " + regAlloc.registerToString(reg) + ", " + std::to_string(offset) + "(sp)");
    }
    int bytesToRestore = savedCallerRegisters.size() * 4;
    emit("    addi sp, sp, " + std::to_string(bytesToRestore));
    savedCallerRegisters.clear();
    currentCallerSaveStackOffset = 0;
}

Register CodeGenerator::generateExpression(ASTNode* node) {
    if (!node) return Register::ZERO;
    
    if (node->type == "AddExpr" || node->type == "MulExpr" || 
        node->type == "RelExpr" || node->type == "LAndExpr" || 
        node->type == "LOrExpr") {
        return generateBinaryOp(node);
    } else if (node->type == "UnaryExpr") {
        return generateUnaryOp(node);
    } else if (node->type == "VarRef") {
        return generateVarRef(node);
    } else if (node->type == "NumLiteral") {
        return generateNumLiteral(node);
    } else if (node->type == "FuncCall") {
        // 递归和普通函数调用统一处理
        return generateFuncCall(node);
    }
    
    return Register::ZERO;
}

Register CodeGenerator::generateBinaryOp(ASTNode* node) {
    if (!node || node->children.size() < 2 || !node->children[0] || !node->children[1]) return Register::ZERO;
    // 严格先左后右
    Register leftReg = generateExpression(node->children[0].get());
    Register rightReg = generateExpression(node->children[1].get());
    if (leftReg == rightReg) {
        Register tmpReg = regAlloc.allocateRegister();
        emit("    mv " + regAlloc.registerToString(tmpReg) + ", " + regAlloc.registerToString(rightReg));
        rightReg = tmpReg;
    }
    Register resultReg = regAlloc.allocateRegister();
    std::string op = node->value;
    if (op == "+") {
        emit("    add " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "-") {
        emit("    sub " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "*") {
        emit("    mul " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "/") {
        emit("    div " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "%") {
        emit("    rem " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "<") {
        emit("    slt " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == ">") {
        emit("    sgt " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
    } else if (op == "<=") {
        emit("    sgt " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    seqz " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(resultReg));
    } else if (op == ">=") {
        emit("    slt " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    seqz " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(resultReg));
    } else if (op == "==") {
        emit("    sub " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    seqz " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(resultReg));
    } else if (op == "!=") {
        emit("    sub " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    snez " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(resultReg));
    } else if (op == "||") {
        // Convert operands to boolean 0/1 then OR
        emit("    snez " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(leftReg));
        emit("    snez " + regAlloc.registerToString(rightReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    or " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    snez " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(resultReg));
    } else if (op == "&&") {
        // Convert operands to boolean 0/1 then AND
        emit("    snez " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(leftReg));
        emit("    snez " + regAlloc.registerToString(rightReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    and " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(leftReg) + ", " + regAlloc.registerToString(rightReg));
        emit("    snez " + regAlloc.registerToString(resultReg) + ", " + regAlloc.registerToString(resultReg));
    }
    regAlloc.freeRegister(leftReg);
    regAlloc.freeRegister(rightReg);
    return resultReg;
}

Register CodeGenerator::generateUnaryOp(ASTNode* node) {
    if (!node || node->children.size() < 1 || !node->children[0]) return Register::ZERO;
    Register operandReg = generateExpression(node->children[0].get());
    Register resultReg = regAlloc.allocateRegister();
    
    std::string op = node->value;
    
    if (op == "+") {
        emit("    mv " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(operandReg));
    } else if (op == "-") {
        emit("    neg " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(operandReg));
    } else if (op == "!") {
        emit("    seqz " + regAlloc.registerToString(resultReg) + ", " + 
             regAlloc.registerToString(operandReg));
    }
    
    regAlloc.freeRegister(operandReg);
    return resultReg;
}

// 在generateVarRef中加调试输出
Register CodeGenerator::generateVarRef(ASTNode* node) {
    if (!node) return Register::ZERO;
    std::string varName = node->value;
    Register reg = regAlloc.allocateRegister();
    int offset = findVariableOffset(varName);
    if (offset == -114514) {
        std::cerr << "[ERROR] generateVarRef: variable '" << varName << "' not found!" << std::endl;
        emit("    li " + regAlloc.registerToString(reg) + ", 0");
    } else if (offset >= 10000 && offset < 10008) {
        // 理论上不会再出现；如果出现说明前面忘记 spill 形参
        int regIdx = offset - 10000;
        std::string regName = regAlloc.registerToString(static_cast<Register>((int)Register::A0 + regIdx));
        emit("    mv " + regAlloc.registerToString(reg) + ", " + regName);
    } else if (offset >= 0) {
        emit("    lw " + regAlloc.registerToString(reg) + ", " + std::to_string(offset) + "(s0)");
    } else if (offset != -114514) {
        emit("    lw " + regAlloc.registerToString(reg) + ", " + std::to_string(offset) + "(s0)");
    } else {
        std::cerr << "[ERROR] generateVarRef: variable '" << varName << "' not found!" << std::endl;
        emit("    li " + regAlloc.registerToString(reg) + ", 0");
    }
    return reg;
}

Register CodeGenerator::generateNumLiteral(ASTNode* node) {
    int value = std::stoi(node->value);
    Register reg = regAlloc.allocateRegister();
    
    if (value >= -2048 && value <= 2047) {
        emit("    li " + regAlloc.registerToString(reg) + ", " + std::to_string(value));
    } else {
        int hi20 = ((value + 0x800) >> 12) & 0xFFFFF;
        int lo12 = value - (hi20 << 12);
        emit("    lui " + regAlloc.registerToString(reg) + ", " + std::to_string(hi20));
        emit("    addi " + regAlloc.registerToString(reg) + ", " +
             regAlloc.registerToString(reg) + ", " + std::to_string(lo12));
    }
    
    return reg;
}

Register CodeGenerator::generateFuncCall(ASTNode* node) {
    std::string funcName = node->value;
    saveRegisters();
    size_t paramCount = node->children.size();
    std::vector<Register> argRegs;
    
    // 先计算所有参数的值
    for (size_t i = 0; i < paramCount; ++i) {
        Register argReg = generateExpression(node->children[i].get());
        argRegs.push_back(argReg);
    }
    
    // 计算需要压栈的参数数量
    size_t stackParamCount = (paramCount > 8) ? (paramCount - 8) : 0;
    
    // 右到左push超过8个参数到栈上
    for (int i = (int)paramCount - 1; i >= 8; --i) {
        pushToStack(argRegs[i]);
    }
    
    // 左到右放前8个参数到a0~a7，从右到左设置避免覆盖
    for (int i = std::min((int)paramCount - 1, 7); i >= 0; --i) {
        Register paramReg = static_cast<Register>((int)Register::A0 + i);
        emit("    mv " + regAlloc.registerToString(paramReg) + ", " + regAlloc.registerToString(argRegs[i]));
        regAlloc.allocateSpecificRegister(paramReg);
    }
    
    emit("    call " + funcName);
    
    // 恢复栈指针，平衡栈
    if (stackParamCount > 0) {
        int stackOffset = stackParamCount * 4; // 每个参数占4字节
        emit("    addi sp, sp, " + std::to_string(stackOffset));
    }
    
    restoreRegisters();
    
    // 释放参数寄存器
    for (size_t i = 0; i < paramCount && i < 8; ++i) {
        Register paramReg = static_cast<Register>((int)Register::A0 + i);
        regAlloc.freeRegister(paramReg);
    }
    
    // 释放参数值寄存器
    for (Register reg : argRegs) {
        regAlloc.freeRegister(reg);
    }
    
    Register retReg = regAlloc.allocateRegister();
    emit("    mv " + regAlloc.registerToString(retReg) + ", a0");
    return retReg;
}

// 重建作用域栈的辅助函数
void CodeGenerator::rebuildScopeStack(ASTNode* node, std::stack<Scope>& tempScopeStack, const std::set<std::string>& paramNameSet) {
    if (!node) return;
    
    // 为嵌套代码块创建新作用域
    if (node->type == "Block") {
        int baseOffset = -4;
        if (!tempScopeStack.empty()) baseOffset = tempScopeStack.top().nextOffset;
        tempScopeStack.push(Scope(baseOffset));
    }
    
    if (node->type == "DeclAssignStmt") {
        std::string varName = node->children[1]->value;
        if (paramNameSet.count(varName)) return; // 形参，跳过
        
        if (!tempScopeStack.empty()) {
            auto& vo = tempScopeStack.top().varOffsets;
            auto it = vo.find(varName);
            if (it != vo.end() && (it->second >= 10000 && it->second < 10008 || it->second >= 0)) {
                // 是形参，跳过
                return;
            } else {
                int offset = tempScopeStack.top().nextOffset;
                tempScopeStack.top().nextOffset = offset - 4; // 递减偏移量（负值）
                tempScopeStack.top().varOffsets[varName] = offset;
                std::cerr << "[rebuild] var: " << varName << " allocated at offset: " << offset << std::endl;
            }
        }
    }
    
    for (const auto& child : node->children) {
        if (child) rebuildScopeStack(child.get(), tempScopeStack, paramNameSet);
    }
    
    // 不再退出嵌套代码块的作用域，保持作用域栈完整
    // if (node->type == "Block") {
    //     if (!tempScopeStack.empty()) {
    //         tempScopeStack.pop();
    //     }
    // }
}

void CodeGenerator::collectVarOffsets(ASTNode* node, const std::set<std::string>& paramNameSet) {
    if (!node) return;
    
    // 为嵌套代码块创建新作用域
    if (node->type == "Block") {
        enterScope();
    }
    
    if (node->type == "DeclAssignStmt") {
        std::string varName = node->children[1]->value;
        std::cerr << "[collect] var: " << varName << ", isParam: " << (paramNameSet.count(varName) ? "true" : "false") << std::endl;
        if (paramNameSet.count(varName)) return; // 形参，跳过
        // 如果是形参（已分配且为reg_a0~reg_a7或正偏移），跳过
        if (!scopeStack.empty()) {
            auto& vo = scopeStack.top().varOffsets;
            auto it = vo.find(varName);
            if (it != vo.end() && (it->second >= 10000 && it->second < 10008 || it->second >= 0)) {
                // 是形参，跳过
                std::cerr << "[collect] var: " << varName << " is param, skipping" << std::endl;
            } else {
                int offset = scopeStack.top().nextOffset;
                scopeStack.top().nextOffset = offset - 4; // 递减偏移量（负值）
                scopeStack.top().varOffsets[varName] = offset;
                allVarOffsets.push_back(offset);
                if (offset < minNextOffset) minNextOffset = offset;
                std::cerr << "[collect] var: " << varName << " allocated at offset: " << offset << std::endl;
            }
        }
    }
    
    for (const auto& child : node->children) {
        if (child) collectVarOffsets(child.get(), paramNameSet);
    }
    
    // 退出嵌套代码块的作用域
    if (node->type == "Block") {
        exitScope();
    }
}

void CodeGenerator::generate(ASTNode* root) {
    if (!root || root->type != "CompUnit") return;
    emitComment("Data section");
    emit(".data");
    emit("");
    emitComment("Text section");
    emit(".text");
    emit(".globl main");
    bool hasMain = false;
    for (const auto& child : root->children) {
        if (child && child->type == "FuncDef" && child->children.size() > 1 && child->children[1] && child->children[1]->type == "ID") {
            std::string funcName = child->children[1]->value;
            if (funcName == "main") hasMain = true;
            currentFunction = funcName;
            
            // 清空之前的状态
            globalVarOffsets.clear();
            nextOffset = 0;
            minNextOffset = 0;
            allVarOffsets.clear();
            currentFunctionParams.clear();
            currentFunctionParamInfo.clear();
            currentFunctionParamOrder.clear();
            
            // 创建函数作用域
            while (!scopeStack.empty()) scopeStack.pop();
            enterScope();
            
            // 收集形参名
            std::set<std::string> paramNameSet;
            if (child->children.size() > 2 && child->children[2] && child->children[2]->type == "ParamList") {
                int paramIdx = 0;
                int stackParamOffset = 0;
                for (const auto& param : child->children[2]->children) {
                    if (!param || param->children.size() < 1) continue;
                    std::string paramName = param->children[0]->value;
                    paramNameSet.insert(paramName);
                    if (paramIdx < 8) {
                        scopeStack.top().varOffsets[paramName] = 10000 + paramIdx; // 特殊标记reg_a0~reg_a7
                    } else {
                        scopeStack.top().varOffsets[paramName] = stackParamOffset; // 正偏移量：0, 4, 8...
                        stackParamOffset += 4;
                    }
                    paramIdx++;
                }
            }
            
            ASTNode* blockNode = nullptr;
            if (child->children.size() > 3 && child->children[3]) blockNode = child->children[3].get();
            
            // 收集变量偏移量
            if (blockNode) collectVarOffsets(blockNode, paramNameSet);
            
            // 计算栈帧大小
            int minOffset = 0;
            for (int off : allVarOffsets) if (off < minOffset) minOffset = off;
            int localVarSize = abs(minOffset); // 使用最小偏移量的绝对值
            int savedRegsSize = 52;
            int frameSize = ((savedRegsSize + localVarSize + 15) / 16) * 16;
            currentFrameSize = frameSize;
            std::cerr << "[frame] func: " << funcName << ", frameSize: " << frameSize << ", localVarSize: " << localVarSize << ", minOffset: " << minOffset << std::endl;
            
            // 创建变量偏移量映射表
            std::unordered_map<std::string, int> varOffsetMap;
            
            // 添加参数到映射表
            if (child->children.size() > 2 && child->children[2] && child->children[2]->type == "ParamList") {
                int paramIdx = 0;
                int stackParamOffset = 0;
                for (const auto& param : child->children[2]->children) {
                    if (!param || param->children.size() < 1) continue;
                    std::string paramName = param->children[0]->value;
                    if (paramIdx < 8) {
                        int offset = scopeStack.top().nextOffset;
                        scopeStack.top().nextOffset = offset - 4;
                        varOffsetMap[paramName] = offset;
                    } else {
                        varOffsetMap[paramName] = stackParamOffset;
                        stackParamOffset += 4;
                    }
                    paramIdx++;
                }
            }
            
            // 重新构建作用域栈并收集变量偏移量
            while (!scopeStack.empty()) scopeStack.pop();
            enterScope();
            
            // 重新分配参数
            if (child->children.size() > 2 && child->children[2] && child->children[2]->type == "ParamList") {
                int paramIdx = 0;
                int stackParamOffset = 0;
                for (const auto& param : child->children[2]->children) {
                    if (!param || param->children.size() < 1) continue;
                    std::string paramName = param->children[0]->value;
                    if (paramIdx < 8) {
                        int offset = scopeStack.top().nextOffset;
                        scopeStack.top().nextOffset = offset - 4;
                        scopeStack.top().varOffsets[paramName] = offset;
                    } else {
                        scopeStack.top().varOffsets[paramName] = stackParamOffset;
                        stackParamOffset += 4;
                    }
                    paramIdx++;
                }
            }
            
            // 重新分配变量偏移量
            if (blockNode) {
                // 重新构建作用域栈
                std::stack<Scope> tempScopeStack;
                tempScopeStack.push(scopeStack.top()); // 复制函数作用域
                
                // 递归重建作用域栈
                rebuildScopeStack(blockNode, tempScopeStack, paramNameSet);
                
                // 将重建的作用域栈复制回来
                scopeStack = tempScopeStack;
                
                // 将变量偏移量添加到映射表
                std::stack<Scope> tempStack = tempScopeStack;
                while (!tempStack.empty()) {
                    const Scope& scope = tempStack.top();
                    for (const auto& pair : scope.varOffsets) {
                        if (pair.second < 10000 && pair.second < 0) { // 只添加局部变量
                            varOffsetMap[pair.first] = pair.second;
                        }
                    }
                    tempStack.pop();
                }
            } else {
                std::cerr << "[debug] blockNode is null!" << std::endl;
            }
            
            // 将变量偏移量映射表存储为全局变量，供代码生成使用
            globalVarOffsets = varOffsetMap;
            
            // 重新初始化作用域栈用于代码生成
            while (!scopeStack.empty()) scopeStack.pop();
            enterScope();
            
            // 重新分配参数到作用域栈
            if (child->children.size() > 2 && child->children[2] && child->children[2]->type == "ParamList") {
                int paramIdx = 0;
                int stackParamOffset = 0;
                for (const auto& param : child->children[2]->children) {
                    if (!param || param->children.size() < 1) continue;
                    std::string paramName = param->children[0]->value;
                    if (paramIdx < 8) {
                        int offset = scopeStack.top().nextOffset;
                        scopeStack.top().nextOffset = offset - 4;
                        scopeStack.top().varOffsets[paramName] = offset;
                    } else {
                        scopeStack.top().varOffsets[paramName] = stackParamOffset;
                        stackParamOffset += 4;
                    }
                    paramIdx++;
                }
            }
            
            emitFunctionPrologue(funcName, frameSize);

            // --- 将最多 8 个寄存器传递的形参保存到栈帧局部空间 ---
            if (child->children.size() > 2 && child->children[2] && child->children[2]->type == "ParamList") {
                int paramIdx = 0;
                for (const auto& param : child->children[2]->children) {
                    if (!param || param->children.size() < 1) continue;
                    std::string paramName = param->children[0]->value;
                    if (paramIdx < 8) {
                        // 为该形参分配一个新的负偏移（像普通局部变量一样）
                        int offset = scopeStack.top().nextOffset;
                        scopeStack.top().nextOffset = offset - 4;
                        scopeStack.top().varOffsets[paramName] = offset;
                        globalVarOffsets[paramName] = offset;

                        Register aReg = static_cast<Register>((int)Register::A0 + paramIdx);
                        emit("    sw " + regAlloc.registerToString(aReg) + ", " + std::to_string(offset) + "(s0)");
                    }
                    paramIdx++;
                }
            }
            
            bool hasReturn = false;
            if (blockNode) generateBlock(blockNode, hasReturn, frameSize, paramNameSet);
            if (!hasReturn) emitFunctionEpilogue(frameSize);
        }
    }
    if (!hasMain) {
        std::cerr << "Error: No main function defined!" << std::endl;
    }
}

void CodeGenerator::generateBlock(ASTNode* node, bool& hasReturn, int frameSize, const std::set<std::string>& paramNameSet) {
    if (!node) return;
    
    // 为块语句创建新作用域
    enterScope();
    
    // 将当前块中的变量添加到作用域栈
    for (const auto& child : node->children) {
        if (child && child->type == "DeclAssignStmt") {
            std::string varName = child->children[1]->value;
            // 从globalVarOffsets中找到对应的偏移量
            auto it = globalVarOffsets.find(varName);
            if (it != globalVarOffsets.end()) {
                // 使用当前作用域的偏移量，而不是全局偏移量
                int currentOffset = scopeStack.top().nextOffset;
                scopeStack.top().nextOffset = currentOffset - 4;
                scopeStack.top().varOffsets[varName] = currentOffset;
            }
        }
    }
    
    for (const auto& child : node->children) {
        if (child) generateStatement(child.get(), hasReturn, frameSize, paramNameSet);
    }
    
    // 退出作用域
    exitScope();
}

void CodeGenerator::generateStatement(ASTNode* node, bool& hasReturn, int frameSize, const std::set<std::string>& paramNameSet) {
    if (!node) return;
    // 移除这个检查，让所有语句都被处理
    // if (hasReturn) return;
    
    if (node->type == "AssignStmt") {
        generateAssignStmt(node);
    } else if (node->type == "DeclAssignStmt") {
        generateDeclAssignStmt(node, paramNameSet);
    } else if (node->type == "IfStmt") {
        generateIfStmt(node, hasReturn, frameSize, paramNameSet);
    } else if (node->type == "IfElseStmt") {
        generateIfElseStmt(node, hasReturn, frameSize, paramNameSet);
    } else if (node->type == "WhileStmt") {
        generateWhileStmt(node, hasReturn, frameSize, paramNameSet);
    } else if (node->type == "BreakStmt") {
        generateBreakStmt(node);
    } else if (node->type == "ContinueStmt") {
        generateContinueStmt(node);
    } else if (node->type == "ReturnStmt") {
        generateReturnStmt(node, hasReturn, frameSize);
    } else if (node->type == "Block") {
        generateBlock(node, hasReturn, frameSize, paramNameSet);
    } else if (node->type == "ExprStmt") {
        // 处理表达式语句，如函数调用
        if (node->children.size() > 0 && node->children[0]) {
            if (node->children[0]->type == "FuncCall") {
                ASTNode* funcCall = node->children[0].get();
                generateFuncCall(funcCall);
            } else {
                generateExpression(node->children[0].get());
            }
        }
    }
}

// 在generateAssignStmt和generateDeclAssignStmt等所有变量写入前加调试输出
void CodeGenerator::generateAssignStmt(ASTNode* node) {
    if (!node || node->children.size() < 2 || !node->children[0] || !node->children[1]) return;
    std::string varName = node->children[0]->value;
    Register valueReg = generateExpression(node->children[1].get());
    int offset = findVariableOffset(varName);
    if (offset == -114514) {
        std::cerr << "[ERROR] generateAssignStmt: variable '" << varName << "' not declared!" << std::endl;
    }
    std::cerr << "[sw-assign] var: " << varName << ", offset: " << offset << ", frameSize: " << currentFrameSize << std::endl;
    emit("    sw " + regAlloc.registerToString(valueReg) + ", " + std::to_string(offset) + "(s0)");
    regAlloc.freeRegister(valueReg);

}

// 在generateDeclAssignStmt中加调试输出
void CodeGenerator::generateDeclAssignStmt(ASTNode* node, const std::set<std::string>& paramNameSet) {
    if (!node || node->children.size() < 3 || !node->children[1] || !node->children[2]) return;
    std::string varName = node->children[1]->value;
    if (paramNameSet.count(varName)) return; // 形参，跳过
    // 如果是形参（已分配且为reg_a0~reg_a7或正偏移），直接return
    if (!scopeStack.empty()) {
        auto& vo = scopeStack.top().varOffsets;
        auto it = vo.find(varName);
        if (it != vo.end() && (it->second >= 10000 && it->second < 10008 || it->second >= 0)) {
            return;
        }
    }
    // 不再在这里分配偏移量，因为已经在collectVarOffsets中分配过了
    Register valueReg = generateExpression(node->children[2].get());
    int realOffset = findVariableOffset(varName);
    if (realOffset == -114514) {
        std::cerr << "[ERROR] generateDeclAssignStmt: variable '" << varName << "' not declared!" << std::endl;
    }
    std::cerr << "[sw-decl] var: " << varName << ", offset: " << realOffset << ", frameSize: " << currentFrameSize << std::endl;
    emit("    sw " + regAlloc.registerToString(valueReg) + ", " + std::to_string(realOffset) + "(s0)");
    regAlloc.freeRegister(valueReg);

}

void CodeGenerator::generateIfStmt(ASTNode* node, bool& hasReturn, int frameSize, const std::set<std::string>& paramNameSet) {
    if (node->children.size() < 2) return;
    std::string endLabel = generateLabel("if_end");

    generateCondition(node->children[0].get(), "", endLabel);

    bool thenReturn = false;
    if (node->children[1]->type == "Block") {
        // 对于块语句，不再进入新作用域，因为作用域栈已经重建
        generateStatement(node->children[1].get(), thenReturn, frameSize, paramNameSet);
    } else {
        generateStatement(node->children[1].get(), thenReturn, frameSize, paramNameSet);
    }

    emit(endLabel + ":");
    hasReturn = thenReturn;
}

void CodeGenerator::generateIfElseStmt(ASTNode* node, bool& hasReturn, int frameSize, const std::set<std::string>& paramNameSet) {
    if (node->children.size() < 3) return;
    std::string elseLabel = generateLabel("else");
    std::string endLabel = generateLabel("if_end");

    generateCondition(node->children[0].get(), "", elseLabel);

    bool thenReturn = false;
    if (node->children[1]->type == "Block") {
        // 对于块语句，不再进入新作用域，因为作用域栈已经重建
        generateStatement(node->children[1].get(), thenReturn, frameSize, paramNameSet);
    } else {
        generateStatement(node->children[1].get(), thenReturn, frameSize, paramNameSet);
    }

    if (!thenReturn) {
        emit("    j " + endLabel);
    }

    emit(elseLabel + ":");

    bool elseReturn = false;
    if (node->children[2]->type == "Block") {
        // 对于块语句，不再进入新作用域，因为作用域栈已经重建
        generateStatement(node->children[2].get(), elseReturn, frameSize, paramNameSet);
    } else {
        generateStatement(node->children[2].get(), elseReturn, frameSize, paramNameSet);
    }

    if (!elseReturn) {
        emit(endLabel + ":");
    }

    hasReturn = thenReturn && elseReturn;
}

void CodeGenerator::generateWhileStmt(ASTNode* node, bool& hasReturn, int frameSize, const std::set<std::string>& paramNameSet) {
    if (node->children.size() < 2) return;
    std::string loopLabel = generateLabel("while_loop");
    std::string endLabel = generateLabel("while_end");

    continueLabels.push(loopLabel);
    breakLabels.push(endLabel);

    emit(loopLabel + ":");
    generateCondition(node->children[0].get(), "", endLabel);

    bool bodyReturn = false;
    if (node->children[1]->type == "Block") {
        // 对于块语句，不再进入新作用域，因为作用域栈已经重建
        generateStatement(node->children[1].get(), bodyReturn, frameSize, paramNameSet);
    } else {
        generateStatement(node->children[1].get(), bodyReturn, frameSize, paramNameSet);
    }

    emit("    j " + loopLabel);
    emit(endLabel + ":");

    continueLabels.pop();
    breakLabels.pop();
}

void CodeGenerator::generateBreakStmt(ASTNode* node) {
    if (!breakLabels.empty()) {
        emit("    j " + breakLabels.top());
    }
}

void CodeGenerator::generateContinueStmt(ASTNode* node) {
    if (!continueLabels.empty()) {
        emit("    j " + continueLabels.top());
    }
}

void CodeGenerator::generateReturnStmt(ASTNode* node, bool& hasReturn, int frameSize) {
    hasReturn = true;
    if (!node) {
        // 空return，a0不变
    } else if (node->value == "with_value" && node->children.size() > 0 && node->children[0]) {
        ASTNode* exprNode = node->children[0].get();
        Register valueReg = generateExpression(exprNode);
        emit("    mv a0, " + regAlloc.registerToString(valueReg));
        regAlloc.freeRegister(valueReg);
    }
    emitFunctionEpilogue(frameSize);
}

void CodeGenerator::generateCondition(ASTNode* node, const std::string& trueLabel, const std::string& falseLabel) {
    if (!node) return;
    // 若为逻辑或/与，需要生成短路控制流以保持语义
    if (node->type == "LOrExpr" || node->type == "LAndExpr") {
        generateShortCircuit(node, trueLabel, falseLabel);
        return;
    }

    // 其余表达式直接求值再判断
    Register resultReg = generateExpression(node);
    emit("    beqz " + regAlloc.registerToString(resultReg) + ", " + falseLabel);
    if (!trueLabel.empty()) {
        emit("    j " + trueLabel);
    }
    regAlloc.freeRegister(resultReg);
}

void CodeGenerator::generateShortCircuit(ASTNode* node, const std::string& trueLabel, const std::string& falseLabel) {
    if (!node) return;

    if (node->type == "LOrExpr") {
        // Logical OR short-circuit: if the left side is true, skip evaluating RHS.
        std::string rightLabel = generateLabel("or_right");
        bool needEndLabel = trueLabel.empty();
        std::string endLabel = needEndLabel ? generateLabel("or_end") : trueLabel;

        // left ? goto endLabel : goto rightLabel
        generateCondition(node->children[0].get(), endLabel, rightLabel);

        // Evaluate RHS only if left side was false
        emit(rightLabel + ":");
        generateCondition(node->children[1].get(), endLabel, falseLabel);

        if (needEndLabel) {
            emit(endLabel + ":");
        }
    } else if (node->type == "LAndExpr") {
        // Logical AND short-circuit: if the left side is false, skip evaluating RHS.
        std::string rightLabel = generateLabel("and_right");
        bool needEndLabel = falseLabel.empty();
        std::string endLabel = needEndLabel ? generateLabel("and_end") : falseLabel;

        // left ? goto rightLabel : goto endLabel
        generateCondition(node->children[0].get(), rightLabel, endLabel);

        // Evaluate RHS only if left side was true
        emit(rightLabel + ":");
        generateCondition(node->children[1].get(), trueLabel, endLabel);

        if (needEndLabel) {
            emit(endLabel + ":");
        }
    }
}

// 修正pushToStack，彻底移除对scopeStack.top().nextOffset的递减，参数压栈和临时空间分配只操作sp
void CodeGenerator::pushToStack(Register reg) {
    emit("    addi sp, sp, -4");
    emit("    sw " + regAlloc.registerToString(reg) + ", 0(sp)");
}

Register CodeGenerator::popFromStack() {
    Register reg = regAlloc.allocateRegister();
    emit("    lw " + regAlloc.registerToString(reg) + ", 0(sp)");
    emit("    addi sp, sp, 4");
    return reg;
}

int CodeGenerator::getVariableOffset(const std::string& varName) {
    // 这个函数现在委托给作用域感知的版本
    return getVariableOffsetInScope(varName);
}

int CodeGenerator::countLocalVars(ASTNode* node) {
    if (!node) return 0;
    int count = 0;
    if (node->type == "DeclAssignStmt") {
        count++;
    } else if (node->type == "AssignStmt") {
        // 检查是否是新的变量赋值（不是已声明的变量）
        if (node->children.size() > 0 && node->children[0]) {
            std::string varName = node->children[0]->value;
            // 检查是否是参数（非main函数中的a,b,n,x,y）
            bool isParameter = false;
            if (currentFunction != "main") {
                if (varName == "a" || varName == "b" || varName == "n" || varName == "x" || varName == "y") {
                    isParameter = true;
                }
            }
            if (!isParameter) {
                count++;
            }
        }
    }
    for (const auto& child : node->children) {
        if (child) count += countLocalVars(child.get());
    }
    return count;
}

// 变量分配辅助函数
// 删除preallocateVariables相关实现
// 统计所有局部变量的最大负偏移
int CodeGenerator::getMinVarOffset() {
    int minOffset = 0;
    std::stack<Scope> tmpStack = scopeStack;
    while (!tmpStack.empty()) {
        const Scope& sc = tmpStack.top();
        for (const auto& pair : sc.varOffsets) {
            if (pair.second < minOffset) minOffset = pair.second;
        }
        tmpStack.pop();
    }
    for (const auto& pair : globalVarOffsets) {
        if (pair.second < minOffset) minOffset = pair.second;
    }
    return minOffset;
}

void CodeGenerator::enterScope() {
    int baseOffset = -4; // 局部变量从-4开始（负偏移量）
    if (!scopeStack.empty()) baseOffset = scopeStack.top().nextOffset;
    scopeStack.push(Scope(baseOffset));
    std::cerr << "[scope-enter] stack size: " << scopeStack.size() << ", baseOffset: " << baseOffset << std::endl;
}

void CodeGenerator::exitScope() {
    if (!scopeStack.empty()) {
        std::cerr << "[scope-exit] stack size: " << scopeStack.size() << ", nextOffset: " << scopeStack.top().nextOffset << std::endl;
        scopeStack.pop();
    }
}

// 修正：禁止变量访问时分配新offset，只允许preallocateVariables分配
int CodeGenerator::getVariableOffsetInScope(const std::string& varName) {
    if (scopeStack.empty()) {
        // 如果没有作用域栈，使用全局偏移量表（兼容性）
        if (globalVarOffsets.find(varName) == globalVarOffsets.end()) {
            // 变量未声明，直接返回特殊值
            std::cerr << "[ERROR] getVariableOffsetInScope: variable '" << varName << "' not declared!" << std::endl;
            return -114514; // 特殊错误码
        }
        return globalVarOffsets[varName];
    }
    // 检查当前作用域是否已经声明了这个变量
    Scope& currentScope = scopeStack.top();
    if (currentScope.varOffsets.find(varName) != currentScope.varOffsets.end()) {
        // 当前作用域已经声明了这个变量，直接返回其偏移量
        return currentScope.varOffsets[varName];
    }
    // 变量未声明，直接返回特殊值
    std::cerr << "[ERROR] getVariableOffsetInScope: variable '" << varName << "' not declared in scope!" << std::endl;
    return -114514;
}

// 修正：变量访问一律用findVariableOffset，找不到就报错
int CodeGenerator::findVariableOffset(const std::string& varName) {
    // 首先检查作用域栈（从栈顶开始，实现变量遮蔽）
    std::stack<Scope> tmpStack = scopeStack;
    while (!tmpStack.empty()) {
        auto& varOffsets = tmpStack.top().varOffsets;
        auto it = varOffsets.find(varName);
        if (it != varOffsets.end()) {
            // 找到变量，返回其偏移量
            return it->second;
        }
        tmpStack.pop();
    }
    
    // 然后检查全局变量偏移量表
    auto globalIt = globalVarOffsets.find(varName);
    if (globalIt != globalVarOffsets.end()) {
        return globalIt->second;
    }
    
    return -114514;
}

bool CodeGenerator::isVariableInCurrentScope(const std::string& varName) {
    if (scopeStack.empty()) {
        return globalVarOffsets.find(varName) != globalVarOffsets.end();
    }
    return scopeStack.top().varOffsets.find(varName) != scopeStack.top().varOffsets.end();
}

void CodeGenerator::printCode() const {
    // 直接输出到标准输出，不保存到内部向量
    // 这样 ./compiler test.c 就等同于 ./compiler < test.c > test.s
    for (const auto& line : code) {
        std::cout << line << std::endl;
    }
}

void CodeGenerator::saveToFile(const std::string& filename) const {
    std::ofstream file(filename);
    if (file.is_open()) {
        for (const auto& line : code) {
            file << line << std::endl;
        }
        file.close();
    }
}

// 优化方法（简化实现）
void CodeGenerator::optimizeCommonSubexpressions() {
    // TODO: 实现公共子表达式消除
}

void CodeGenerator::optimizeConstantFolding() {
    // TODO: 实现常量折叠
}

void CodeGenerator::optimizeRegisterAllocation() {
    // TODO: 实现更高级的寄存器分配
}